import os
from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceBgeEmbeddings
from langchain.text_splitter import CharacterTextSplitter

# --- 配置 ---
KNOWLEDGE_BASE_DIR = "data_sample" # 知识库目录
FAISS_INDEX_DIR = "faiss_index"       # 向量数据库保存目录
EMBEDDING_MODEL = "BAAI/bge-small-zh" # 和你程序中一致的嵌入模型

def create_vector_store():
    """从知识库目录加载文档，创建并保存FAISS向量数据库。"""
    print(f"--- 开始从 '{KNOWLEDGE_BASE_DIR}' 加载文档 ---")

    # 加载所有 .txt 文档
    documents = []
    for filename in os.listdir(KNOWLEDGE_BASE_DIR):
        if filename.endswith(".txt"):
            file_path = os.path.join(KNOWLEDGE_BASE_DIR, filename)
            loader = TextLoader(file_path, encoding='utf-8')
            documents.extend(loader.load())

    if not documents:
        print("错误：在知识库目录中没有找到任何 .txt 文件。")
        return

    print(f"成功加载 {len(documents)} 篇文档。")

    # 分割文档
    text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
    docs = text_splitter.split_documents(documents)
    print(f"文档被分割成 {len(docs)} 个片段。")

    # 加载嵌入模型
    print(f"--- 正在加载嵌入模型: {EMBEDDING_MODEL} ---")
    # HuggingFaceBgeEmbeddings 默认使用 CUDA (如果可用)
    model_kwargs = {'device': 'cuda'}
    encode_kwargs = {'normalize_embeddings': True}
    embeddings = HuggingFaceBgeEmbeddings(
        model_name=EMBEDDING_MODEL,
        model_kwargs=model_kwargs,
        encode_kwargs=encode_kwargs
    )
    print("嵌入模型加载完成。")

    # 创建并构建 FAISS 向量数据库
    print("--- 正在创建 FAISS 向量数据库... ---")
    vector_store = FAISS.from_documents(docs, embeddings)

    # 保存到本地
    vector_store.save_local(FAISS_INDEX_DIR)
    print(f"--- 向量数据库成功创建并保存于 '{FAISS_INDEX_DIR}' 文件夹中！ ---")

if __name__ == "__main__":
    create_vector_store()